/* -*-c++-*- SemiGlobalMatching - Copyright (C) 2020.
* Author	: Ethan Li <ethan.li.whu@gmail.com>
* Describe	: main
*/

#include "stdafx.h"
#include "SemiGlobalMatching.h"

// opencv library
#include <opencv2/opencv.hpp>
#ifdef _DEBUG
#pragma comment(lib,"opencv_world320d.lib")
#else
#pragma comment(lib,"opencv_world320.lib")
#endif

void disp2Depth(cv::Mat& dispMap, cv::Mat& depthMap, cv::Mat& K);

/**
 * \brief 
 * \param argv 2~4
 * \param argc argc[1]:左影像路径 argc[2]: 右影像路径 argc[3]: 最小视差[可选，默认0] argc[4]: 最大视差[可选，默认64]
 * \param eg. ..\Data\cone\im2.png ..\Data\cone\im6.png 0 64 
 * \return 
 */
int main(int argv,char** argc)
{
    if(argv < 3) {
        std::cout << "请指定左右摄像头照片路径!" << std::endl;
        return -1;
    }

    //···············································································//
    // 读取影像
    std::string path_left = argc[1];
    std::string path_right = argc[2];
   
    cv::Mat img_left = cv::imread(path_left, cv::IMREAD_GRAYSCALE);
    cv::Mat img_right = cv::imread(path_right, cv::IMREAD_GRAYSCALE);

    if (img_left.data == nullptr || img_right.data == nullptr) {
        std::cout << "读取照片失败!" << std::endl;
        return -1;
    }
    if (img_left.rows != img_right.rows || img_left.cols != img_right.cols) {
        std::cout << "左右照片尺寸不一致!" << std::endl;
        return -1;
    }


    //···············································································//
    const sint32 width = static_cast<uint32>(img_left.cols);
    const sint32 height = static_cast<uint32>(img_right.rows);

	// 左右影像的灰度数据
    auto bytes_left = new unsigned char[width*height];
    auto bytes_right = new unsigned char[width*height];
	for (int i = 0; i < height;i++) {
		for (int j = 0; j < width;j++) {
			bytes_left[i*width + j] = img_left.at<unsigned char>(i, j);
			bytes_right[i*width + j] = img_right.at<unsigned char>(i, j);
		}
	}


    // SGM匹配参数设计
    SemiGlobalMatching::SGMOption sgm_option;
    // 聚合路径数
	sgm_option.num_paths = 8;
	// 候选视差范围
    sgm_option.min_disparity = argv < 4 ? 0 : atoi(argc[3]);
    sgm_option.max_disparity = argv < 5 ? 64 : atoi(argc[4]);
    // 一致性检查
    sgm_option.is_check_lr = true;
    sgm_option.lrcheck_thres = 1.0f;
    // 唯一性约束
	sgm_option.is_check_unique = true;
    sgm_option.uniqueness_ratio = 0.99;
    // 剔除小连通区
	sgm_option.is_remove_speckles = true;
    sgm_option.min_speckle_aera = 30;
    // 惩罚项P1、P2
	sgm_option.p1 = 10;
    sgm_option.p2_init = 150;
	// 视差图填充
	sgm_option.is_fill_holes = true;
    
	// 定义SGM匹配类实例
    SemiGlobalMatching sgm;

    //···············································································//
    // 初始化
    if(!sgm.Initialize(width, height, sgm_option)) {
        std::cout << "SGM初始化失败！" << std::endl;
        return -2;
    }


    //···············································································//
    // 匹配
    auto disparity = new float32[uint32(width * height)]();
    if(!sgm.Match(bytes_left,bytes_right,disparity)) {
        std::cout << "SGM匹配失败!" << std::endl;
        return -2;
    }

	//···············································································//
    // 显示视差图
    cv::Mat disp_mat = cv::Mat(height, width, CV_8UC1);
    float min_disp = width, max_disp = -width;
	for (sint32 i = 0; i < height; i++) {
        for (sint32 j = 0; j < width; j++) {
            const float32 disp = disparity[i * width + j];
            if (disp != Invalid_Float) {
                min_disp = std::min(min_disp, disp);
                max_disp = std::max(max_disp, disp);
            }
        }
    }
    for (sint32 i = 0; i < height; i++) {
        for (sint32 j = 0; j < width; j++) {
            const float32 disp = disparity[i * width + j];
            if (disp == Invalid_Float) {
                disp_mat.data[i * width + j] = 0;
            }
            else {
                disp_mat.data[i * width + j] = static_cast<uchar>((disp - min_disp) / (max_disp - min_disp) * 255);
            }
        }
    }

    cv::imshow("disparity", disp_mat);
	cv::Mat disp_color;
	applyColorMap(disp_mat, disp_color, cv::COLORMAP_JET);
	cv::imshow("disparity-color", disp_color);

	// 保存结果
    std::string disp_map_path = argc[2]; disp_map_path += ".d.bmp";
    std::string disp_color_map_path = argc[2]; disp_color_map_path += ".c.bmp";
    cv::imwrite(disp_map_path, disp_mat);
    cv::imwrite(disp_color_map_path, disp_color);

    //--- add by Leon ---
    //读取标定数据
    cv::Mat camMtx[2], distCoeffs[2];
    cv::Mat R, T;
    cv::Mat leftCamMtx, rightCamMtx;

    cv::FileStorage fs("intrinsics.yml", cv::FileStorage::READ);
    if (fs.isOpened())
    {
        fs["M1"] >> camMtx[0];
        fs["D1"] >> distCoeffs[0];
        fs["M2"] >> camMtx[1];
        fs["D2"] >> distCoeffs[1];
        fs["R"]  >> R;
        fs["T"]  >> T;
        fs.release();
    }

    cv::transpose(camMtx[0], leftCamMtx);
//    cv::transpose(camMtx[1], rightCamMtx);

    float fx = leftCamMtx.at<float>(0, 0);

    // 转换成深度图，并保存结果
    cv::Mat depthMap(height, width, CV_16UC1);
    cv::Mat leftCam;

    disp2Depth(disp_mat, depthMap, leftCamMtx);

    float baseline = -T.at<float>(0, 0);    // 镜头之间的距离
    std::cout << "baseline:" << baseline << std::endl;

    uchar* dispData = (uchar*)disp_mat.data;
    ushort* depthData = (ushort*)depthMap.data;
    
    float minDepth = 50000.;
    int posX = 0xdead;
    int posY = 0xdead;

    // 遍历所有点位距离，找出距离最近的点位位置
    for (sint32 i = 0; i < height; i++) {
        for (sint32 j = 0; j < width; j++) {
            int id = i * width + j;
            const float32 disp = disparity[id];
            if (disp != Invalid_Float) {
                float depth = (float)fx * baseline / ((float)dispData[id]);
                depthData[id] = depth;

                if (depth < minDepth) {
                    minDepth = depth;
                    posX = j;
                    posY = i;
                }
            }
        }
    }

    if (posX != 0xdead && posY != 0xdead){
        std::cout << "pos_x:" << posX << "pos_y:" << posY << std::endl;
    }
    else{
        std::cout << "error! 没有找到相应位置!" << std::endl;
    }
    //--- Leon end ---

	cv::waitKey(0);
    
    //--- add by Leon  ---
    depthMap.release();
    //--- Leon end ---

    //.........................................................................
    // 释放内存
    delete[] disparity; 
	disparity = nullptr;
	delete[] bytes_left;
	bytes_left = nullptr;
	delete[] bytes_right;
	bytes_right = nullptr;

    system("pause");
	return 0;
}

/*
函数作用：视差图转深度图
输入：
　　dispMap ----视差图，8位单通道，CV_8UC1
　　K       ----内参矩阵，float类型
输出：
　　depthMap ----深度图，16位无符号单通道，CV_16UC1
*/
void disp2Depth(cv::Mat& dispMap, cv::Mat& depthMap, cv::Mat& K)
{
    int type = dispMap.type();

    float fx = K.at<float>(0, 0);
    float fy = K.at<float>(1, 1);
    float cx = K.at<float>(0, 2);
    float cy = K.at<float>(1, 2);
    float baseline = 38; //基线距离38mm

    if (type == CV_8U)
    {
        const float PI = 3.14159265358;
        int height = dispMap.rows;
        int width = dispMap.cols;

        uchar* dispData = (uchar*)dispMap.data;
        ushort* depthData = (ushort*)depthMap.data;
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                int id = i * width + j;
                if (!dispData[id])  continue;  //防止0除
                depthData[id] = ushort((float)fx * baseline / ((float)dispData[id]));
            }
        }
    }
    else
    {
        std::cout << "please confirm dispImg's type!" << std::endl;
        cv::waitKey(0);
    }
}